package com.dc.project.util;

import com.dc.project.property.ProjectVolcengineProperty;
import com.dc.project.service.UserVolcengineService;
import com.own.component.common.login.util.SessionUserUtil;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import lombok.extern.slf4j.Slf4j;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * VolcengineChatUtil
 *
 * @author chenxueli
 * @date 2025-03-07 16:00:00
 */
@Slf4j
@Component
public class VolcengineChatUtil {

    private static ArkService SERVICE = null;

    @Resource
    private ProjectVolcengineProperty projectVolcengineProperty;

    @Resource
    private UserVolcengineService userVolcengineService;

    /**
     * 获取调用的service
     *
     * @return ArkService
     */
    private ArkService getArkService() {
        // 根据登录用户的大模型配置去获取对应的参数
        var loginUser = SessionUserUtil.getLoginUser(false);
        if (loginUser.isLogin()) {
            var userVolcengine = userVolcengineService.getByUserIdForCache(loginUser.userId());
            if (userVolcengine != null) {
                return ArkService.builder()
                        .dispatcher(new Dispatcher())
                        .connectionPool(new ConnectionPool(5, 1, TimeUnit.SECONDS))
                        .baseUrl(userVolcengine.getBaseUrl())
                        .apiKey(userVolcengine.getApiKey())
                        .build();
            }
        }
        // 如果用户自己没有配置大模型配置，则使用默认的配置（后期需要移除该逻辑）
        if (SERVICE == null) {
            var connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
            var dispatcher = new Dispatcher();
            SERVICE = ArkService.builder()
                    .dispatcher(dispatcher)
                    .connectionPool(connectionPool)
                    .baseUrl(projectVolcengineProperty.getBaseUrl())
                    .apiKey(projectVolcengineProperty.getArkApiKey())
                    .build();
        }
        return SERVICE;
    }

    /**
     * 调用接口
     *
     * @param systemPrompt   系统提示词
     * @param userPromptList 用户内容
     * @return 结果
     */
    public String invoke(String systemPrompt, List<String> userPromptList) {
        try {
            var service = getArkService();
            final var messages = new ArrayList<ChatMessage>();
            final var systemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content(systemPrompt).build();
            messages.add(systemMessage);
            // 循环添加用户内容
            userPromptList.forEach(prompt -> {
                final var userMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(prompt).build();
                messages.add(userMessage);
            });

            var chatCompletionRequest = ChatCompletionRequest.builder().model(getModel()).messages(messages).build();

            var result = service.createChatCompletion(chatCompletionRequest).getChoices().stream().map(choice -> choice.getMessage().getContent()).map(Object::toString).collect(Collectors.joining(""));
            service.shutdownExecutor();
            return result;
        } catch (Exception e) {
            return "";
        }
    }

    /**
     * 调用接口
     *
     * @param systemPrompt 系统提示词
     * @param userPrompt   用户内容
     * @return 结果
     */
    public String invoke(String systemPrompt, String... userPrompt) {
        return invoke(systemPrompt, List.of(userPrompt));
    }

    /**
     * 获取模型
     *
     * @return 模型id
     */
    private String getModel() {
        // 根据登录用户的大模型配置去获取对应的参数
        var loginUser = SessionUserUtil.getLoginUser(false);
        if (loginUser.isLogin()) {
            var userVolcengine = userVolcengineService.getByUserIdForCache(loginUser.userId());
            if (userVolcengine != null) {
                return userVolcengine.getModelId();
            }
        }
        return projectVolcengineProperty.getModel();
    }

}
